"""
Equal Weighting Strategy
等权重策略实现
"""

from typing import Dict, List, Optional
import numpy as np
from ml_lib.core import Tensor
from .base import WeightingStrategy


class EqualWeighting(WeightingStrategy):
    """
    等权重策略
    
    最简单的权重策略，所有任务的权重都相等。
    """
    
    def __init__(self, task_names: List[str], **kwargs):
        """
        初始化等权重策略
        
        Args:
            task_names: 任务名称列表
        """
        super().__init__(task_names, **kwargs)
        
    def compute_weights(self, losses: Dict[str, Tensor], 
                       gradients: Optional[Dict[str, List[np.ndarray]]] = None,
                       **kwargs) -> Dict[str, float]:
        """
        计算任务权重（等权重）
        
        Args:
            losses: 各任务的损失值字典
            gradients: 各任务的梯度字典（未使用）
            **kwargs: 其他参数（未使用）
            
        Returns:
            任务权重字典，所有权重都为1.0
        """
        return {name: 1.0 for name in self.task_names}
    
    def __repr__(self):
        return f"EqualWeighting(num_tasks={self.num_tasks})" 